# -*-coding:utf-8-*-
import cv2
import numpy as np


class Segmentation:
    def __init__(self, img):
        if isinstance(img, str):
            self.img = cv2.imread(img)
        elif isinstance(img, np.ndarray):
            self.img = img
        self.b, self.g, self.r = None, None, None
        self.select_channel = None
        self.thre = 140

    @staticmethod
    def cv_show(name, img):
        cv2.imshow(name, img)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    @staticmethod
    def median_blur(img, ksize=9):
        median = cv2.medianBlur(img, ksize)  # en
        result = cv2.absdiff(median, img)
        return result

    @staticmethod
    def edge_connect(img, k=20):
        gray = img / 255.0  # 像素值0-1之间

        # sobel算子分别求出gx，gy
        gx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
        gy = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
        mag, ang = cv2.cartToPolar(gx, gy, angleInDegrees=1)  # 得到梯度幅度和梯度角度阵列
        g = np.zeros(gray.shape)  # g与图片大小相同

        # 行扫描，间隔k时，进行填充，填充值为1
        def edge_connection(img, size, k):
            for i in range(size):
                Yi = np.where(img[i, :] > 0)
                if len(Yi[0]) >= 10:  # 可调整
                    for j in range(0, len(Yi[0]) - 1):
                        if Yi[0][j + 1] - Yi[0][j] <= k:
                            img[i, Yi[0][j]:Yi[0][j + 1]] = 1
            return img

        # 选取边缘，提取边缘坐标，将g中相应坐标像素值设为1
        X, Y = np.where((mag > np.max(mag) * 0.3) & (ang >= 0) & (ang <= 90))
        g[X, Y] = 1

        # 边缘连接，此过程只涉及水平，垂直边缘连接，不同角度边缘只需旋转相应角度即可
        g = edge_connection(g, gray.shape[0], k)
        g = cv2.rotate(g, 0)
        g = edge_connection(g, gray.shape[1], k)
        g = cv2.rotate(g, 2)
        return g

    @staticmethod
    def dilate_iter(img, epoch=200):
        '''
        img:grayscale, binary image
        膨胀迭代，封闭区域填充
        :return: img
        '''

        thresh_not = cv2.bitwise_not(img)  # 二值图像的补集
        # cv_show('thresh_not', thresh_not)

        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))  # 3×3结构元

        # 构建阵列F，并将thresh_not边界值写入F
        F = np.zeros(img.shape, np.uint8)
        F[:, 0] = thresh_not[:, 0]
        F[:, -1] = thresh_not[:, -1]
        F[0, :] = thresh_not[0, :]
        F[-1, :] = thresh_not[-1, :]

        # 循环迭代，对F进行膨胀操作，结果与thresh_not执行and操作
        for i in range(epoch):
            F_dilation = cv2.dilate(F, kernel, iterations=1)
            F = cv2.bitwise_and(F_dilation, thresh_not)
        #     res = np.hstack((thresh_not, F_dilation, F))
        #     cv_show('res', res)

        result = cv2.bitwise_not(F)  # 对结果执行not

        return result

    @staticmethod
    def ed_denoising(img, ksize):
        '''
        膨胀 腐蚀去除噪点

        :return: a binary denoising image
        '''
        kernel = np.ones((ksize, ksize), np.uint8)
        erosion = cv2.erode(img, kernel, iterations=1)
        midres2 = cv2.dilate(erosion, kernel, iterations=1)
        return midres2

    @staticmethod
    def caculate_pixels(root):
        color_black = np.where(root == 0)[0].shape[0]
        color_white = np.where(root == 255)[0].shape[0]

        pixel_sum = root.shape[0] * root.shape[1]
        return color_black/pixel_sum

    @staticmethod
    def adaptive_select_channel(img):
        b, g, r = cv2.split(img)
        # cv_show('merge', np.hstack((b, g, r)))
        max = 0
        min = 255
        for i in range(3):
            if img[:, :, i].mean() > max:
                max = img[:, :, i].mean()
                channel_max = i
            if img[:, :, i].mean() < min:
                min = img[:, :, i].mean()
                channel_min = i
        channel = [j for j in range(3) if (j != channel_max) & (j != channel_min)]
        # print(channel)
        if channel[0] == 0:
            return channel, b
        elif channel[0] == 1:
            return channel, g
        else:
            return channel, r

    @staticmethod
    def get_variable_name(var):
        loc = locals()
        for k, v in loc.items():
            if loc[k] == var:
                return k

    def main_pipline(self):
        edgec = 1 # 是否要边缘连接
        # 中值滤波
        img = self.img
        self.b, self.g, self.r = cv2.split(img)
        # cv_show('name', img)
        median2 = self.median_blur(img, 5)  # en

        # 二值化 腐蚀膨胀 去噪
        res2_gray = cv2.cvtColor(median2, cv2.COLOR_BGR2GRAY)
        ret, midres = cv2.threshold(res2_gray, 85, 255, cv2.THRESH_BINARY)   # 划痕图1
        # cv_show('可取', midres)

        if edgec:
            midres = self.ed_denoising(midres, 2)

            # 边缘连接
            g = self.edge_connect(midres)    # 划痕图2
            # cv_show('a', midres)
            # cv_show('g', g)

        # 通道二值化
        if isinstance(self.select_channel, np.ndarray):
            channel = self.select_channel
            # self.select_channel = None
        else:
            index, channel = self.adaptive_select_channel(img)  # 取通道
        thre = self.thre
        ret, thrb = cv2.threshold(channel, thre, 255, cv2.THRESH_BINARY)    # 调整阈值, b:default = 60, g: 220
        # ret, thrb_r = cv2.threshold(r, 100, 255, cv2.THRESH_BINARY)

        thrb = self.ed_denoising(thrb, 2)
        # thrb_r = ed_denoising(thrb_r, 2)
        # cv_show('thrb', thrb)

        # 膨胀迭代
        if self.caculate_pixels(thrb) <= 0.5:
            thrb = cv2.bitwise_not(thrb)

        result = self.dilate_iter(thrb, 400)

        if self.caculate_pixels(result) >= 0.5:
            result = cv2.bitwise_not(result)
        # cv_show('mask', result)

        # if os.path.exists('./segresult' + '/' + name):
        #     pass
        # else:
        #     home = os.getcwd()
        #     os.mkdir(os.path.join(home, 'segresult') + '/' + name)
        # cv2.imwrite('./segresult/%s/mask.png' % name, result)

        # 凹坑及划痕
        # y = np.expand_dims(rea, -1)
        y = np.stack((result, result, result), axis=2)
        p = img * (y / 255).astype(np.uint8)    # 划痕凹坑图
        # cv_show('p ', p)
        # cv2.imwrite('haha', p)
        return result, p

    def seg_detect(self, img, index=None, bitnot1=False, bitnot2=False):
        edgec = 1  # 是否要边缘连接
        # 中值滤波
        b, g, r = cv2.split(img)
        self.b, self.g, self.r = 'b', 'g', 'r'
        median2 = self.median_blur(img, 5)  # en

        # 二值化 腐蚀膨胀 去噪
        res2_gray = cv2.cvtColor(median2, cv2.COLOR_BGR2GRAY)
        ret, midres = cv2.threshold(res2_gray, 85, 255, cv2.THRESH_BINARY)  # 划痕图1
        # cv_show('可取', midres)

        if edgec:
            midres = self.ed_denoising(midres, 2)

            # 边缘连接
            edg_connect = self.edge_connect(midres)  # 划痕图2
            # cv_show('a', midres)
            # cv_show('g', g)

        # b通道二值化
        if isinstance(self.select_channel, str):
            if self.select_channel == 'b':
                channel = b
            elif self.select_channel == 'g':
                channel = g
            else:
                channel = r
        else:
            if index:
                if index[0] == 0:
                    channel = b
                elif index[0] == 1:
                    channel = g
                else:
                    channel = r
            else:
                index, channel = self.adaptive_select_channel(img)  # 取通道
        thre = self.thre
        ret, thrb = cv2.threshold(channel, thre, 255, cv2.THRESH_BINARY)  # 调整阈值, b:default = 60, g: 220
        # ret, thrb_r = cv2.threshold(r, 100, 255, cv2.THRESH_BINARY)

        thrb = self.ed_denoising(thrb, 2)
        # cv_show('thrb', thrb)

        # 膨胀迭代
        if thrb.shape == (480, 640):
            if self.caculate_pixels(thrb) <= 0.5:
                thrb = cv2.bitwise_not(thrb)
                bitnot1 = True

            result = self.dilate_iter(thrb, 400)

            if self.caculate_pixels(result) >= 0.5:
                result = cv2.bitwise_not(result)
                bitnot2 = True
        else:
            if bitnot1:
                thrb = cv2.bitwise_not(thrb)

            result = self.dilate_iter(thrb, 400)

            if bitnot2:
                result = cv2.bitwise_not(result)

        # 凹坑及划痕
        # y = np.expand_dims(rea, -1)
        y = np.stack((result, result, result), axis=2)
        p = img * (y / 255)  # 划痕凹坑图
        # cv_show('p ', p)
        # cv2.imwrite('./segresult/%s/result.png' % name, p)
        return result, p, index, bitnot1, bitnot2


if __name__ == '__main__':
    img = "../images/18.png"
    main_pipline(img)